import tensorflow as tf


def parser(record):
    keys_to_features = {
        'image': tf.FixedLenFeature((), tf.string),
        'label': tf.FixedLenFeature((), tf.int64)
    }
    parsed = tf.io.parse_single_example(record, keys_to_features)
    image_ = tf.decode_raw(parsed['image'], tf.uint8)
    image_ = tf.cast(image_, tf.float32)
    label_ = tf.cast(parsed['label'], tf.int32)
    return image_, label_


if __name__ == '__main__':
    dataset = tf.data.TFRecordDataset('data/train.tfrecord')
    dataset = dataset.map(parser)
    dataset = dataset.batch(5)
    dataset = dataset.repeat()
    iterator = dataset.make_one_shot_iterator()
    feature, label = iterator.get_next()
    with tf.Session() as sess:
        print(sess.run(label))
        print(sess.run(feature))
